import itertools
import matplotlib.pyplot as plt
from util import iter_kwds

__all__ = [ 'ts_plot' ]

def ts_plot(ts, *args, **kwds):
    '''Plot time series object 'ts'. If the times series is not unidimensional,
    idx can be an int or a list of ints specifying the dimensions that should
    be plotted. '''

    kwdsit = iter_kwds(kwds)
    idx = kwds.pop('idx', None)
    if idx is None:
        if ts.data_shape:
            idx_lst = itertools.product(*map(range, ts.data_shape))
        else:
            idx_lst = [ None ]
    elif isinstance(idx, list):
        idx_lst = idx
    else:
        idx_lst = [ idx ]

    for idx in idx_lst:
        ts_i = ts.view_dim(idx)
        X = ts_i.times
        Y = ts_i.data
        plt.plot(X, Y, *args, **kwdsit.next())

if __name__ == '__main__':
    import doctest
    doctest.testmod(optionflags=doctest.REPORT_ONLY_FIRST_FAILURE)
